from __future__ import print_function, division
import time
#Importing relevant modules and packages
import numpy as np
import matplotlib.pyplot as plt
import random
from RichardsModel_1D import *
from Parameters import *
from scipy import integrate
from tqdm import tqdm
from casadi import *



p=Loam() #Soil parameters

#Space and time parameters
Nr,Nt,Nz,dr,dt,dz,Np,NN,Nx,Nw,Ny,Nv,Hz=circular_parameters()
DeltaT,Nsim=time_parameters()

N = 50 # Number of particles

#Irrigation amount
# uu=np.zeros(360*40)
# for i in range(360*40):
#     if i < 360:
#         uu[i]=-1e-05 #Moisture is supplied only at the first time instant
#     else:
#         uu[i]=0
    # uu[0*i]=-1e-05

u=np.zeros(10*120)
uu=np.zeros(Nsim)
for i in range(10*120):
    if i < 50:
        u[i]=-2e-06 #Moisture is supplied only at the first time instant
    else:
        u[i]=0
uu=np.tile(u,(1,120))
uu=np.transpose(uu)


Kc=0.88*np.ones(Nsim)      #Crop coefficient
Kc[2700:5400]=1.08
Kc[5400:8100]=0.88
Kc[8100:Nsim]=1.08
Kc1=1.8*np.ones(Nsim)

Et=1.389e-08*np.ones(Nsim) #Reference evapotranspiration
Et[2700:5400]=1.489e-08
Et[5400:8100]=1.389e-08
Et[8100:Nsim]=1.489e-08
Et1=1.289e-08*np.ones(Nsim) #Reference evapotranspiration


#Initial condition [pressure head]
x0=-0.8*np.ones(Nz)

#Numpy array to store the simulation results [pressure head]
xx=np.zeros((Nsim+1,Nz)) #true state
xx1=np.zeros((Nsim+1,Nz)) #open loop
xx_cur=np.zeros((Nsim+1,Nz)) #filter
xx_pre = np.zeros((Nsim+1,Nz)) #predict

xx[0,:]=x0
# xx1[0,:]=1.1*x0
xx1[0,:]=1.1*x0
xx_cur[0,:]=1.1*x0

# array for filtering
x_out = np.zeros((Nsim+1,Nz))
x_est = np.zeros((Nsim+1,Nz))
x_est_out = np.zeros((Nsim+1,Nz))
x_out[0,:] = x0  #the actual output vector for measurement values.
x_est[0,:] = xx_cur[0,:] # time by time output of the particle filters estimate

#setting P/Q/R
Sigma_p=10
Sigma_w=0.1
Sigma_v=550

P_pre=np.diag(Sigma_p*np.ones((Nz,)))
Q=np.diag(Sigma_w*np.ones((Nz,)))
R=np.diag(Sigma_v*np.ones((Nz,)))

# setting noise
www1 = 4e-9 # process noise
vv1 = np.array([8e-7,8e-7,8e-7,8e-7,8e-7,8e-7,8e-7,8e-7,8e-7,8e-7,8e-7,8e-7,8e-7,8e-7,8e-7,8e-7]) # measurement noise
# vv1 = 8e-07*np.ones([1,16])
# vv1 = np.array([1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6,1e-6]) # measurement noise
vv2 = np.diag(vv1)
vvv1 = 4e-9 # particle noise
wwww1 = 4e-9

np.random.seed(1)
ww=0*np.random.normal(0,np.sqrt(0.01),[Nsim,Nz]) #If no model uncertainty is present, ww is multiplied by 0
# np.random.seed(2)
www = np.random.normal(0,np.sqrt(www1),[Nsim+1,Nz])
# np.random.seed(3)
# vv = np.random.normal(0,np.sqrt(vv1),[Nsim+1,Nz]) #measurement noise
vv = np.random.randn(Nsim+1,Nz)*vv1
# np.random.seed(4)
vvv = np.random.normal(0,np.sqrt(vvv1),[N,Nz]) # initial particle covariance
# np.random.seed(5)
wwww = np.random.normal(0,np.sqrt(wwww1),[Nsim,Nz])


#Numpy array to store the simulation results [volumetric moisture content]
yy=np.zeros((Nsim+1,Nz)) # measurements for ture states 
yy1=np.zeros((Nsim+1,Nz)) # measurements for open loop
yy2=np.zeros((Nsim+1,Nz)) 
yy[0,:]=getOutputs1D_allnodes(x0).full().ravel()+vv[0,:]
yy1[0,:]=getOutputs1D_allnodes(xx1[0,:]).full().ravel()+vv[0,:]

#unknown inputs
a1=0*np.ones((Nsim+1,Nz))
# a=0.000009*np.ones((Nsim+1,Nz))#constant UIs
a=0.00000*np.ones((Nsim+1,Nz))
# for i in range (Nsim+1):
#     a[i,:]=a[i,:]+(1e-06)*np.cos(0.0005*i)

# a[:,11]=0.00005
# a[:,12]=0.0001
# a[:,13]=0.0002
# a[:,14]=0.0004
# a[:,15]=0.0003
a_est=0.00000*np.ones((Nsim+1,Nz))
# a_est[:,11]=0.00002
# a_est[:,12]=0.00005
# a_est[:,13]=0.00015
# a_est[:,14]=0.00035
# a_est[:,15]=0.00025

# gamma=[0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.018,0.018,0.018,0.018] #6*40
# gamma=[0.15,0.15,0.15,0.15,0.15,0.15,0.15,0.15,0.15,0.15,0.15,0.15,0.025,0.025,0.025,0.025]
# gamma=[0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.0035,0.0035,0.0035,0.0035]
# gamma=[0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.018,0.0035,0.0035,0.0035,0.0035]


 

e0 = time.time()
# NOW THE FUNCTION IS COMPILED, RE-TIME IT EXECUTING FROM CACHE


    
for i in tqdm(range(1, Nsim+1)):
    # Generate observed data (x bias)
    xx[i,:]= (DM(ODE_discrete(tuple(xx[i-1]), uu[i-1],ww[i-1],Kc[i-1],Et[i-1],DeltaT)).full().ravel())+a[i-1,:]+www[i-1]
    
    yy[i,:]=getOutputs1D_allnodes(xx[i,:]).full().ravel()+vv[i,:]

for i in tqdm(range(1, Nsim+1)):
    # open loop
    xx1[i,:]= (DM(ODE_discrete(tuple(xx1[i-1]), uu[i-1],ww[i-1],Kc1[i-1],Et1[i-1],DeltaT)).full().ravel())+a_est[i-1,:]+www[i-1]
    
    yy1[i,:]=getOutputs1D_allnodes(xx1[i,:]).full().ravel()+vv[i,:]

# for i in tqdm(range(1, Nsim+1)):
#     # Data with noise without UIs
#     ode2[i-1,:]= RichardsModel_1D(tuple(xx1[i-1]), uu[i-1],ww[i-1],a1[i-1],Kc[i-1],Et[i-1],p)
#     xx1[i,:]=xx1[i-1,:]+DeltaT*ode2[i-1,:]+wwww[i-1]
#     yy1[i,:]=getOutputs1D_allnodes(xx1[i,:],p).ravel()+vv[i,:]
    
# x_est_out = np.loadtxt('xooutest.txt')

t= list(range(Nsim+1))
plt.figure(1)
for i in range (Nz):
    plt.subplot(4, 4, i+1)
    # plt.plot(t,xx[:,i],'r',t,xx1[:,i],'b',t,x_est_out[:,i],'g')
    # plt.plot(t,xx[:,i],'r',t,xx1[:,i],'b',t,xx2[:,i],'g')
    plt.plot(t,xx[:,i],'r',t,xx1[:,i],'b')
    # plt.plot(t,xx[:,i],'r',t,xx1[:,i])
    plt.legend(['True value','open loop','without UIs'])

plt.figure(2)
for i in range (Nz):
    plt.subplot(4, 4, i+1)
    # plt.plot(t,xx[:,i],'r',t,xx1[:,i],'b',t,x_est_out[:,i],'g')
    # plt.plot(t,xx[:,i],'r',t,xx1[:,i],'b',t,xx2[:,i],'g')
    plt.plot(t,a[:,i],'b')
    # plt.plot(t,xx[:,i],'r',t,xx1[:,i])
    plt.legend(['True value','open loop','without UIs'])
plt.show()
